import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter


class Wangqi(nn.Module):

    # def __init__(self, *args, **kwargs) -> None:
    #     super().__init__(*args, **kwargs)
    #     self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0)
    def __init__(self):
        super(Wangqi,self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0)

    def forward(self, x):
        x = self.conv1(x)
        return x


dataset = torchvision.datasets.CIFAR10("../CIFAR10_dataset", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)
dataloader = DataLoader(dataset, batch_size=64)
wangqi = Wangqi()
Writer = SummaryWriter("logs")
step = 0
for data in dataloader:
    imgs, targets = data
    output = wangqi(imgs)
    Writer.add_images("input", imgs, step)
    output = torch.reshape(output, (-1, 3, 30, 30))
    Writer.add_images("output", output, step)
    step = step + 1
Writer.close()
